import asyncio
from elasticsearch import AsyncElasticsearch
import numpy as np
import re
from typing import List, Dict, Any, Optional

class OptimizedSearch:
    def __init__(self, es: Optional[AsyncElasticsearch] = None, emb_model=None, index_name: str = "jd_file_index_formal"):
        self.es = es
        self.emb_model = emb_model
        self.index_name = index_name
        # 如果没有提供ES客户端或嵌入模型，则需要在使用前设置
        self._fallback_mode = (es is None or emb_model is None)
    
    async def extract_keywords(self, query: str) -> List[str]:
        """从查询中提取关键词"""
        try:
            import jieba
            # 移除标点符号
            query = re.sub(r'[^\w\s]', ' ', query)
            # 分词并过滤停用词
            words = jieba.lcut(query)
            # 过滤短词和停用词
            keywords = [w for w in words if len(w) > 1]
            return keywords[:5]  # 返回最多5个关键词
        except ImportError:
            # 如果没有安装jieba，使用简单的空格分割
            words = query.split()
            return [w for w in words if len(w) > 1][:5]
    
    async def generate_embeddings(self, query: str) -> List[float]:
        """为查询生成嵌入"""
        if not self.emb_model:
            # 如果没有嵌入模型，返回随机向量
            return np.random.rand(1024).tolist()
        
        # 使用嵌入模型生成向量
        try:
            results = await self.emb_model.encode([query])
            return results['dense_vecs'][0].tolist()
        except Exception as e:
            print(f"嵌入生成失败: {e}")
            return np.random.rand(1024).tolist()
    
    async def search(self, query: str, company_id: str = None, limit: int = 20) -> Dict:
        """优化的搜索函数"""
        if self._fallback_mode:
            print("警告：检索系统处于回退模式，某些组件未初始化")
            return {"query": query, "results": [], "total": 0}
        
        # 并行处理查询向量和关键词
        embedding_task = asyncio.create_task(self.generate_embeddings(query))
        keywords_task = asyncio.create_task(self.extract_keywords(query))
        
        # 等待结果
        results = await asyncio.gather(embedding_task, keywords_task)
        query_embedding = results[0]
        keywords = results[1]
        
        # 构建优化的查询
        optimized_query = {
            "size": limit,
            "query": {
                "bool": {
                    "should": [
                        # 多向量字段的相似度查询
                        {
                            "script_score": {
                                "query": {"match_all": {}},
                                "script": {
                                    "source": """
                                        float title_score = cosineSimilarity(params.query_vector, 'title_vector') * 2.0;
                                        float summary_score = cosineSimilarity(params.query_vector, 'summary_text_vector') * 1.5;
                                        float query1_score = cosineSimilarity(params.query_vector, 'query_vector_1') * 1.0;
                                        return title_score + summary_score + query1_score;
                                    """,
                                    "params": {"query_vector": query_embedding}
                                }
                            }
                        },
                        # 关键词匹配
                        {"terms": {"keywords": keywords, "boost": 1.2}},
                        # 全文搜索
                        {"match": {"title": {"query": query, "boost": 2.0}}},
                        {"match": {"summary_text": {"query": query, "boost": 1.5}}},
                        {"match": {"query_1": {"query": query, "boost": 1.0}}},
                        {"match": {"query_2": {"query": query, "boost": 1.0}}},
                        {"match": {"query_3": {"query": query, "boost": 1.0}}}
                    ],
                    "minimum_should_match": 1,
                    "filter": [] if not company_id else [{"term": {"company": company_id}}]
                }
            },
            "_source": ["summary_text", "title", "query_1", "query_2", "query_3", "file_url", "filename", "suffix"]
        }
        
        try:
            # 执行查询
            response = await self.es.search(index=self.index_name, body=optimized_query)
            
            return {
                "query": query,
                "results": response["hits"]["hits"],
                "total": response["hits"]["total"]["value"]
            }
        except Exception as e:
            print(f"ES检索错误: {e}")
            return {
                "query": query,
                "results": [],
                "total": 0,
                "error": str(e)
            }
            
    async def local_search(self, query: str, knowledge_base, threshold: float = 0.2, top_n: int = 5) -> List[Dict]:
        """
        当ES不可用时，使用本地知识库进行检索
        这是一个与现有KnowledgeBase.search兼容的简单接口
        """
        if knowledge_base:
            return knowledge_base.search(query, threshold=threshold, top_n=top_n)
        return [] 